[Stacked PR 5/5] Add GDN backward pass 8k and 64k latency and memory benchmark suite - #5098
[Stacked PR 5/5] Add GDN backward pass 8k and 64k latency and memory benchmark suite#5098Rohan-Bierneni wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a fused analytical Pallas Gated Delta Net (GDN) kernel and backward pipeline, integrating it into the Qwen3 model. Key changes include the addition of the hybrid_bwd_analytical_pipeline and a suite of supporting local GDN kernel files to handle causal Conv1D and Gated Delta Rule operations with cached triangular inverse matrices. Additionally, sublane tiling logic is adjusted in the GMM and TGMM kernels. Feedback on the changes highlights a numerical inconsistency where an epsilon of 1e-12 is used for L2 normalization in the backward pass instead of the 1e-6 used in the forward pass, as well as an opportunity to simplify the sublane alignment assignment in the TGMM kernel.
|
|
||
| if use_qk_norm_in_gdn: | ||
| d_q_scaled = d_q_proj * scale | ||
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12) |
There was a problem hiding this comment.
The epsilon value 1e-12 used here for calculating the norm of q is inconsistent with the value 1e-6 used in the forward pass (l2norm at line 521). For numerical stability and correctness of the gradient, the same epsilon value should be used in both forward and backward computations.
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12) | |
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-6) |
| - q_unit * jnp.sum(d_q_scaled * q_unit, axis=-1, keepdims=True) | ||
| ) / r_q | ||
|
|
||
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12) |
There was a problem hiding this comment.
Similar to the calculation for r_q, the epsilon value 1e-12 used here for r_k is inconsistent with the 1e-6 used in the forward pass. This should be corrected to 1e-6 to ensure numerical consistency.
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12) | |
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-6) |
| common_sublane = min(size_lhs_sublane, size_rhs_sublane) | ||
| size_lhs_sublane = common_sublane | ||
| size_rhs_sublane = common_sublane |
There was a problem hiding this comment.
The logic to ensure size_lhs_sublane and size_rhs_sublane are equal can be simplified into a single line, improving readability.
| common_sublane = min(size_lhs_sublane, size_rhs_sublane) | |
| size_lhs_sublane = common_sublane | |
| size_rhs_sublane = common_sublane | |
| size_lhs_sublane = size_rhs_sublane = min(size_lhs_sublane, size_rhs_sublane) |
83ce15c to
5a5586c
Compare
Rohan-Bierneni
changed the title
5a5586c to
307676e
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Rohan-Bierneni
changed the title
d9b661b to
a127aa0
Compare
Rohan-Bierneni
changed the title
a127aa0 to
aeece03
Compare
…benchmark suite - Add standalone 8k and 64k Ghostfish Cloud TPU latency and memory benchmark suite (gdn_benchmark_test.py). - Provide comparative benchmarking between Pure JAX remat and Decoupled GDN Pallas backward pass kernel across sequence lengths up to 64k.
aeece03 to
56758af
Compare
1 participant
Stacked PR Chain
rbierneni-gdn-1-ci-hygienemainrbierneni-gdn-2-fwd-kernelrbierneni-gdn-1-ci-hygienerbierneni-gdn-3-bwd-kernelrbierneni-gdn-2-fwd-kernelrbierneni-gdn-4-model-integrationrbierneni-gdn-3-bwd-kernelrbierneni-gdnv3-bwdrbierneni-gdn-4-model-integrationDescription
This is PR 5 of 5 (the capstone PR #5098) in the stacked series enabling the Pallas Gated Delta Net (GDN) backward pass kernel in MaxText.
This PR adds the comprehensive Cloud TPU latency and memory benchmark suite comparing the Decoupled GDN Pallas backward pass kernel against Pure JAX remat:
tests/unit/gdn_benchmark_test.py):Files Changed
tests/unit/gdn_benchmark_test.py: Standalone 8k and 64k Cloud TPU latency and memory benchmark suite.Tests
pre-commit run --files ...(all hooks passed).python3 tests/unit/gdn_benchmark_test.pyon Cloud TPU.Checklist